import glob
import random
import numpy as np
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import os
from PIL import Image
import torch

class Multi_Task_Dataset(Dataset):
    def __init__(self, root: str, split: str = "train", transform=None, image_size=256,
                 v_patch_nums=(1, 2, 3, 4, 5, 6, 8, 10, 13, 16), separator=False, **kwargs):
        self.transforms = transform
        self.split = split
        self.image_size = image_size
        self.v_patch_nums = v_patch_nums
        self.separator = separator
        
        # 数据集根目录
        self.root = root
        self.deblur_root = '/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/dataset/gopro'
        self.color_root = '/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/dataset/gopro'
        self.super_root = '/ifs/data/PlatforMax/common/Dataset/DIV2K'
        self.low_light_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Delowlight'
        self.denoise_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Denoise/SIDD'
        self.deraindrop_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Deraindrop'
        self.derain_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Derain'
        self.dehaze_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Dehaze/raw/Haze4K'
        self.desnow_root = '/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Desnow/SRRS-2021'
        self.deblur_color_root = '/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/dataset/gopro'
        # 定义各数据集路径
        self.dataset_roots = {
            'deblur': self.deblur_root,
            'color': self.color_root,
            'low_light': self.low_light_root,
            'super': self.super_root,
            'denoise': self.denoise_root,
            'derain': self.derain_root,
            'deraindrop': self.deraindrop_root,
            'dehaze': self.dehaze_root,
            'desnow': self.desnow_root,
            'deblur_color': self.deblur_color_root
        }
        
        # 加载各数据集的图像对
        self.dataset_pairs = {}
        
        if split == 'train':
            # GoPro数据集（去模糊）
            gopro_blur = sorted(glob.glob(os.path.join(self.dataset_roots['deblur'], split, "input", "*.png"))) #模糊图
            gopro_sharp = sorted(glob.glob(os.path.join(self.dataset_roots['deblur'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['deblur'] = list(zip(gopro_blur, gopro_sharp))
            print(f'Found {len(gopro_blur)} images in GoPro {split} set')
        
            # # DIV2K数据集（超分辨率）
            # div2k_lr = sorted(glob.glob(os.path.join(self.dataset_roots['super'], f"DIV2K_{split}_LR_bicubic", "X2", "*.png"))) #低分辨率图
            # div2k_hr = sorted(glob.glob(os.path.join(self.dataset_roots['super'], f"DIV2K_{split}_HR", "*.png"))) #高分辨率图
            # self.dataset_pairs['super'] = list(zip(div2k_lr, div2k_hr))
            # print(f'Found {len(div2k_lr)} images in DIV2K {split} set')
            
            # 低光图增强
            low_light_input = sorted(glob.glob(os.path.join(self.dataset_roots['low_light'], split, "low", "*.png"))) #模糊图
            low_light_target = sorted(glob.glob(os.path.join(self.dataset_roots['low_light'], split, "high", "*.png"))) #清晰图
            self.dataset_pairs['low_light'] = list(zip(low_light_input, low_light_target))
            print(f'Found {len(low_light_input)} images in LowLight {split} set')
        
            # # SIDD数据集（去噪）
            # sidd_noisy = sorted(glob.glob(os.path.join(self.dataset_roots['denoise'], split, "input_crops", "*.png"))) #噪声图
            # sidd_clean = sorted(glob.glob(os.path.join(self.dataset_roots['denoise'], split, "target_crops", "*.png"))) #清晰图
            # self.dataset_pairs['denoise'] = list(zip(sidd_noisy, sidd_clean))
            # print(f'Found {len(sidd_noisy)} images in SIDD {split} set')
            
            # Deraindrop数据集（去雨滴）
            deraindrop_input = sorted(glob.glob(os.path.join(self.dataset_roots['deraindrop'], split, "rain", "*.png"))) #雨滴图
            deraindrop_target = sorted(glob.glob(os.path.join(self.dataset_roots['deraindrop'], split, "gt", "*.png"))) #清晰图
            self.dataset_pairs['deraindrop'] = list(zip(deraindrop_input, deraindrop_target))
            print(f'Found {len(deraindrop_input)} images in Deraindrop {split} set')
        
            # Rain13k数据集（去雨）
            rain_input = sorted(glob.glob(os.path.join(self.dataset_roots['derain'], split, "Rain13K" , "input", "*.jpg"))) #雨图
            rain_target = sorted(glob.glob(os.path.join(self.dataset_roots['derain'], split, "Rain13K", "target", "*.jpg"))) #清晰图
            self.dataset_pairs['derain'] = list(zip(rain_input, rain_target))
            print(f'Found {len(rain_input)} images in Rain13k {split} set')
        
            # Haze数据集（去雾）
            haze_input = sorted(glob.glob(os.path.join(self.dataset_roots['dehaze'], split, "haze", "*.png"))) #雾图
            haze_target = sorted(glob.glob(os.path.join(self.dataset_roots['dehaze'], split,'gt', "*.png"))) #清晰图
            self.dataset_pairs['dehaze'] = list(zip(haze_input, haze_target))
            print(f'Found {len(haze_input)} images in Haze4K {split} set')
        
            # Snow100K数据集（去雪）
            snow_input = sorted(glob.glob(os.path.join(self.dataset_roots['desnow'], "snow", "*.png"))) #雪图
            snow_target = sorted(glob.glob(os.path.join(self.dataset_roots['desnow'], "gt", "*.jpg"))) #清晰图
            self.dataset_pairs['desnow'] = list(zip(snow_input, snow_target))
            print(f'Found {len(snow_input)} images in Snow100K {split} set')
        
            # Color数据集（上色）
            color_input = sorted(glob.glob(os.path.join(self.dataset_roots['color'], split, "input", "*.png"))) #模糊图
            color_target = sorted(glob.glob(os.path.join(self.dataset_roots['color'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['color'] = list(zip(color_input, color_target))
            print(f'Found {len(color_input)} images in Color {split} set')
        
            # Deblur_Color数据集（灰度图去模糊上色）
            deblur_color_input = sorted(glob.glob(os.path.join(self.dataset_roots['deblur_color'], split, "input", "*.png"))) #模糊图
            deblur_color_target = sorted(glob.glob(os.path.join(self.dataset_roots['deblur_color'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['deblur_color'] = list(zip(deblur_color_input, deblur_color_target))
            print(f'Found {len(deblur_color_input)} images in Deblur_Color {split} set')
            
        else:
            # GoPro数据集（去模糊）
            gopro_blur = sorted(glob.glob(os.path.join(self.dataset_roots['deblur'], split, "input", "*.png"))) #模糊图
            gopro_sharp = sorted(glob.glob(os.path.join(self.dataset_roots['deblur'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['deblur'] = list(zip(gopro_blur, gopro_sharp))
            self.dataset_pairs['deblur_color'] = list(zip(gopro_blur, gopro_sharp))
            print(f'Found {len(gopro_blur)} images in GoPro {split} set')
        
            # # DIV2K数据集（超分辨率）
            # div2k_lr = sorted(glob.glob(os.path.join(self.dataset_roots['super'], "DIV2K_valid_LR_bicubic", "X2", "*.png"))) #低分辨率图
            # div2k_hr = sorted(glob.glob(os.path.join(self.dataset_roots['super'], "DIV2K_valid_HR", "*.png"))) #高分辨率图
            # self.dataset_pairs['super'] = list(zip(div2k_lr, div2k_hr))
            # print(f'Found {len(div2k_lr)} images in DIV2K {split} set')
            
            # 低光图增强
            low_light_input = sorted(glob.glob(os.path.join(self.dataset_roots['low_light'], split, "low", "*.png"))) #模糊图
            low_light_target = sorted(glob.glob(os.path.join(self.dataset_roots['low_light'], split, "high", "*.png"))) #清晰图
            self.dataset_pairs['low_light'] = list(zip(low_light_input, low_light_target))
            print(f'Found {len(low_light_input)} images in LowLight {split} set')
        
            # # SIDD数据集（去噪）
            # sidd_noisy = sorted(glob.glob(os.path.join(self.dataset_roots['denoise'], split, "input_crops", "*.png"))) #噪声图
            # sidd_clean = sorted(glob.glob(os.path.join(self.dataset_roots['denoise'], split, "target_crops", "*.png"))) #清晰图
            # self.dataset_pairs['denoise'] = list(zip(sidd_noisy, sidd_clean))
            # print(f'Found {len(sidd_noisy)} images in SIDD {split} set')
            
            # Deraindrop数据集（去雨滴）
            deraindrop_input = sorted(glob.glob(os.path.join(self.dataset_roots['deraindrop'], split, "rain", "*.jpg"))) #雨滴图
            deraindrop_target = sorted(glob.glob(os.path.join(self.dataset_roots['deraindrop'], split, "gt", "*.jpg"))) #清晰图
            self.dataset_pairs['deraindrop'] = list(zip(deraindrop_input, deraindrop_target))
            print(f'Found {len(deraindrop_input)} images in Deraindrop {split} set')
        
            # Rain13k数据集（去雨）
            rain_input = sorted(glob.glob(os.path.join(self.dataset_roots['derain'], split, "Rain100H" , "input", "*.png"))) #雨图
            rain_target = sorted(glob.glob(os.path.join(self.dataset_roots['derain'], split, "Rain100H", "target", "*.png"))) #清晰图
            self.dataset_pairs['derain'] = list(zip(rain_input, rain_target))
            print(f'Found {len(rain_input)} images in Rain13k {split} set')
        
            # Haze数据集（去雾）
            haze_input = sorted(glob.glob(os.path.join(self.dataset_roots['dehaze'], split, "haze", "*.png"))) #雾图
            haze_target = sorted(glob.glob(os.path.join(self.dataset_roots['dehaze'], split,'gt', "*.png"))) #清晰图
            self.dataset_pairs['dehaze'] = list(zip(haze_input, haze_target))
            print(f'Found {len(haze_input)} images in Haze4K {split} set')
        
            # Snow100K数据集（去雪）
            snow_input = sorted(glob.glob(os.path.join(self.dataset_roots['desnow']+'-val', split , "snow", "*.png"))) #雪图
            snow_target = sorted(glob.glob(os.path.join(self.dataset_roots['desnow']+'-val', split, "gt", "*.jpg"))) #清晰图
            self.dataset_pairs['desnow'] = list(zip(snow_input, snow_target))
            print(f'Found {len(snow_input)} images in Snow100K {split} set')
        
            # Color数据集（上色）
            color_input = sorted(glob.glob(os.path.join(self.dataset_roots['color'], split, "input", "*.png"))) #模糊图
            color_target = sorted(glob.glob(os.path.join(self.dataset_roots['color'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['color'] = list(zip(color_input, color_target))
            print(f'Found {len(color_input)} images in Color {split} set')
        
            # Deblur_Color数据集（灰度图去模糊上色）
            deblur_color_input = sorted(glob.glob(os.path.join(self.dataset_roots['deblur_color'], split, "input", "*.png"))) #模糊图
            deblur_color_target = sorted(glob.glob(os.path.join(self.dataset_roots['deblur_color'], split, "target", "*.png"))) #清晰图
            self.dataset_pairs['deblur_color'] = list(zip(deblur_color_input, deblur_color_target))
            print(f'Found {len(deblur_color_input)} images in Deblur_Color {split} set')
        
        self.default_cls = 0
        
        # 任务ID映射
        self.task_ids = {
            'deblur': 0,          # 去模糊任务
            'color': 1,           # 灰度图上色
            'low_light': 2,      # 低光图增强
            'deraindrop': 3,     # 去雨滴   
            'derain': 4,          # 去雨
            'dehaze': 5,          # 去雾
            'desnow': 6,          # 去雪
            'deblur_color': 7,    # 灰度图去模糊上色
        }
        
        # 计算最小数据集大小
        self.min_dataset_size = min(len(pairs) for pairs in self.dataset_pairs.values())

    def __len__(self):
        return self.min_dataset_size * 8  # 确保每个epoch都能平均采样到每个任务

    def __getitem__(self, index: int):
        # 随机选择任务类型
        task_type = random.choices(
            ['deblur', 'color', 'low_light', 'deraindrop', 'derain', 'dehaze', 'desnow', 'deblur_color'],
            [0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125],
            k=1
        )[0]
        
        # 随机选择图像对
        pair_idx = random.randint(0, len(self.dataset_pairs[task_type]) - 1)
        input_path, target_path = self.dataset_pairs[task_type][pair_idx]
        
        # 读取图像
        try:
            if task_type == 'deblur_color':
                input_img = Image.open(input_path).convert('L').convert('RGB') #deblur的灰度图
                target_img = Image.open(target_path).convert('RGB') #sharp的彩色图
            elif task_type == 'color':
                input_img = Image.open(target_path).convert('L').convert('RGB') #sharp的灰度图
                target_img = Image.open(target_path).convert('RGB') #sharp的彩色图
            else:
                input_img = Image.open(input_path).convert('RGB') # INPUT
                target_img = Image.open(target_path).convert('RGB') # TARGET
        except Exception as e:
            print(f"Error loading images for task {task_type}: {e}")
            print(f"Input path: {input_path}")
            print(f"Target path: {target_path}")
            return self.__getitem__(random.randint(0, len(self) - 1))
        
        # 应用数据增强
        if self.transforms:
            target_img, input_img = self.transforms(target_img, input_img)
        
        ignore_masks = torch.ones((1378,)) if self.separator else torch.ones((1360,))
        ignore_masks_ = torch.ones((1378,)) if self.separator else torch.ones((1360,))
        
        return {
            'image': target_img,
            'mask': input_img,
            'cls': self.default_cls,
            'type': torch.tensor(self.task_ids[task_type]),
            'ignore_mask': ignore_masks,
            'ignore_mask_': ignore_masks_
        }